import torch.nn as nn


class RankLoss(nn.Module):
    def __init__(self):
        super(RankLoss, self).__init__()

    def forward(self, s, p):
        # find the pairs s_i < p_j
        batch, _ = s.shape
        s1 = s.unsqueeze(1)
        s2 = s.unsqueeze(-1)
        s = s1-s2
        s[s < 0] = 0
        s[s > 0] = 1
        p1 = p.unsqueeze(1)
        p2 = p.unsqueeze(-1)
        p = p2 - p1
        p = s*p
        p = 1 - p
        p[p<0] = 0
        return p.sum()/batch

class MaxMinNormal(nn.Module):
    def __init__(self):
        super(MaxMinNormal, self).__init__()

    def forward(self, A):
        A -= A.min(1, keepdim=True)[0]
        A /= A.max(1, keepdim=True)[0]
        return A